package edu.sysu.demo.petstore.dog.key;

import com.google.common.collect.ImmutableList;
import com.google.common.collect.ImmutableMap;
import com.google.common.io.BaseEncoding;
import edu.sysu.demo.petstore.dog.data.dao.KeyDao;
import edu.sysu.demo.petstore.dog.data.entity.Key;
import edu.sysu.demo.petstore.dog.data.entity.Key.KeyType;
import java.security.GeneralSecurityException;
import java.security.KeyPair;
import java.security.KeyPairGenerator;
import java.security.NoSuchAlgorithmException;
import java.security.PrivateKey;
import java.security.Provider;
import java.security.PublicKey;
import java.util.Map;
import java.util.Optional;
import javax.annotation.PostConstruct;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Component;

@Component
public class RSAKeyService implements KeyService {

  private final KeyDao keyDao;

  private final BaseEncoding baseEncoding;

  private final Provider provider;

  private Map<KeyType, KeyConverter> converters;

  @Autowired
  public RSAKeyService(KeyDao keyDao, BaseEncoding baseEncoding, Provider provider) {
    this.keyDao = keyDao;
    this.baseEncoding = baseEncoding;
    this.provider = provider;
  }

  @PostConstruct
  public void initialize() throws GeneralSecurityException {
    converters = ImmutableMap.of(
        KeyType.RSA_PUBLIC, new PublicKeyConverter(baseEncoding, provider),
        KeyType.RSA_PRIVATE, new PrivateKeyConverter(baseEncoding, provider)
    );
  }

  @Override
  public void regenerate(final String name) throws NoSuchAlgorithmException {
    KeyPairGenerator keyPairGenerator = KeyPairGenerator.getInstance("RSA", provider);
    keyPairGenerator.initialize(2048);
    KeyPair keyPair = keyPairGenerator.generateKeyPair();

    PublicKey publicKey = keyPair.getPublic();
    PrivateKey privateKey = keyPair.getPrivate();

    String privateString = baseEncoding.encode(privateKey.getEncoded());
    String publicString = baseEncoding.encode(publicKey.getEncoded());

    Key key1 = Key.builder()
        .content(privateString)
        .name(name)
        .type(KeyType.RSA_PRIVATE)
        .build();

    Key key2 = Key.builder()
        .content(publicString)
        .name(name)
        .type(KeyType.RSA_PUBLIC)
        .build();

    keyDao.saveAll(ImmutableList.of(key1, key2));
  }

  @Override
  public <T extends java.security.Key> T query(String name, KeyType keyType) {

    String encoded = Optional.ofNullable(keyDao.findByNameAndType(name, keyType))
        .map(Key::getContent)
        .orElse(null);

    if (encoded == null) {
      return null;
    }

    try {
      KeyConverter keyConverter = converters.get(keyType);
      return keyConverter.convert(encoded);
    } catch (GeneralSecurityException e) {
      throw new KeyManagementException(e);
    }
  }
}

